package com.alan.boot.mongo.lock;

import org.springframework.data.mongodb.core.FindAndModifyOptions;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;
import org.springframework.stereotype.Service;
import javax.annotation.Resource;
import java.util.List;

/**
 * 基于 mongodb 实现分布式锁
 *
 * mongo 与 redis 的区别：
 * redis 可以 setNx
 * mongo 可以 插入或者 inc 后检查自己是否是第一个，但是不能 setNx
 *
 * mongo 通过 findAndModify 提供了 CAS 的原子操作，可以知道更新前和更新后的值
 * @author yaoyunqing
 * Created on 2020/04/22
 */
@Service
public class DistributionLockService {
    public static final String LOCK_COLLECTION_NAME = "mongoLocks";
    @Resource
    private MongoTemplate mongoTemplate;

    /**
     * 获得锁
     * @param key 锁的键
     * @param expireSeconds 超时时间
     * @return
     */
    public boolean getLock(String key, int expireSeconds) {
        Query keyQuery = new Query();
        keyQuery.addCriteria(Criteria.where("key").is(key));
        List<MongoLock> mongoLocks = mongoTemplate.find(keyQuery, MongoLock.class, LOCK_COLLECTION_NAME);
        // 1. 如果锁已经被持有且未超时，判断获得锁失败
        if (mongoLocks.size() > 0 && mongoLocks.get(0).getExpire() >= System.currentTimeMillis()) {
            return false;
        }
        // 2. 如果锁已经过期了，根据 key + expire 准确释放过期的锁
        if (mongoLocks.size() > 0 && mongoLocks.get(0).getExpire() < System.currentTimeMillis()) {
            releaseLockByKeyAndExpire(key, mongoLocks.get(0).getExpire());
        }
        // 3. 尝试获得锁，注意这个操作即使获得锁失败，也会更新锁的 expire
        Update update = new Update();
        update.inc("value", 1);
        long expire = System.currentTimeMillis() + expireSeconds * 1000;
        update.set("expire", expire);
        // upsert + inc 并获得新值
        FindAndModifyOptions options = FindAndModifyOptions.options().upsert(true).returnNew(true);
        MongoLock mongoLock = mongoTemplate.findAndModify(keyQuery, update, options, MongoLock.class, LOCK_COLLECTION_NAME);
        if (mongoLock.getValue() == 1) {
            return true;
        } else {
            return false;
        }
    }

    /**
     * 根据 锁的键与有效期 释放锁，带上 expire 避免误删锁
     * @param key
     * @param expire
     */
    private void releaseLockByKeyAndExpire(String key, long expire) {
        Query query = new Query();
        query.addCriteria(Criteria.where("key").is(key));
        query.addCriteria(Criteria.where("expire").is(expire));
        mongoTemplate.remove(query, LOCK_COLLECTION_NAME);
    }

    /**
     * 释放指定锁
     * @param key
     */
    public void releaseLock(String key) {
        Query query = new Query();
        query.addCriteria(Criteria.where("key").is(key));
        mongoTemplate.remove(query, LOCK_COLLECTION_NAME);
    }


}
